import pdfplumber

ALLOWED_DOC_EXTENSIONS = ["txt", "pdf", "docx"]
ALLOWED_IMAGE_EXTENSIONS = ["png", "jpeg", "jpg", "gif"]

# 判断文件名是否为允许的文档类型
def allowed_doc_file(filename):
    return filename.rsplit(".", 1)[1].lower() in ALLOWED_DOC_EXTENSIONS

# 提取文档内容，根据不同文件类型分别处理
def extract_doc_content(file_path):
    """
    根据文件类型{txt,pdf,docx} 提取内容
    :param file_path: 文件路径
    :return: 文件字符串内容
    """
    # 获取文件扩展名并转小写
    ext = file_path.rsplit(".")[-1].lower()
    if ext == "txt":
        # 以utf-8编码读取文本内容
        with open(file_path, "r", encoding="utf-8") as f:
            return f.read()
    elif ext == "pdf":
        # 导入pdfplumber
        text = ""
        with pdfplumber.open(file_path) as pdf:
            # 遍历每一页
            for page in pdf.pages:
                for page in pdf.pages:
                    text += page.extract_text() or ""
            # 返回提取的文本内容
            return text
    # 如果是docx文件
    elif ext == "docx":
        # 导入Document类
        from docx import Document
        doc = Document(file_path)
        return "\n".join([para.text for para in doc.paragraphs])
    else:
        raise ValueError(f"不支持的文件类型{ext}")

def allowed_image_file(filename):
    return filename.rsplit(".", 1)[1].lower() in ALLOWED_IMAGE_EXTENSIONS